# -*- coding: utf-8 -*-
# @Time : 2021/11/17 23:19
# @Author : Xiang Qian Xiang Qian
# @Email : qianxljp@126.com
# @File : test.py
# @Project : hrrp_jt
from jittor.utils.pytorch_converter import convert

pytorch_code="""
import torch 
class TripletCELoss(torch.nn.Module):

    def __init__(self, margin=0.1, gama=0.01):
        super(TripletCELoss, self).__init__()
        self.margin = margin
        self.gama = gama
        self.ranking_loss = torch.nn.MarginRankingLoss(margin=margin)

    def forward(self, inputs, targets):

        n0 = inputs.size(0)
        # Compute pairwise distance, replace by the official when merged
        dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n0, n0)
        dist = dist + dist.t()
        dist.addmm_(inputs, inputs.t(), beta=1, alpha=-2)
        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability

        # For each anchor, find the hardest positive and negative
        mask = targets.expand(n0, n0).eq(targets.expand(n0, n0).t())
        dist_ap, dist_an = [], []
        for i in range(n0):
            dist_ap.append(dist[i][mask[i]].max().unsqueeze(0))
            dist_an.append(dist[i][mask[i] == 0].min().unsqueeze(0))
        dist_ap = torch.cat(dist_ap)
        dist_an = torch.cat(dist_an)
        # Compute ranking hinge loss
        y = torch.ones_like(dist_an)
        return self.gama * self.ranking_loss(dist_an, dist_ap, y) + torch.nn.functional.cross_entropy(inputs, targets)
        
        """

jittor_code = convert(pytorch_code)
print(jittor_code)

